package com.example.algorithm.no0337.impl;

import java.util.HashMap;
import java.util.Map;

import com.example.algorithm.no0337.Solution;
import com.example.algorithm.no0337.TreeNode;

/**
 * @author yumuhui <yumuhui@kuaishou.com>
 * Created on 2021-05-13
 */
public class MapSolutionImpl implements Solution {

    // select(o) 表示选择 o 节点的情况下，o 节点的子树上被选择的节点的最大权值和；
    // unSelect(o) 表示不选择 o 节点的情况下，o 节点的子树上被选择的节点的最大权值和；l 和 r 代表 o 的左右孩子。
    Map<TreeNode, Integer> select = new HashMap<>();
    Map<TreeNode, Integer> unSelect = new HashMap<>();

    @Override
    public int rob(TreeNode root) {
        dfs(root);
        return Math.max(select.getOrDefault(root, 0), unSelect.getOrDefault(root, 0));
    }

    public void dfs(TreeNode node) {
        if (node == null) {
            return;
        }
        dfs(node.left);
        dfs(node.right);
        select.put(node, node.val + unSelect.getOrDefault(node.left, 0) + unSelect.getOrDefault(node.right, 0));
        unSelect.put(node, Math.max(select.getOrDefault(node.left, 0), unSelect.getOrDefault(node.left, 0))
                + Math.max(select.getOrDefault(node.right, 0), unSelect.getOrDefault(node.right, 0))
        );
    }

}
